import os
import time
import typing

from FormUi.MainWindowUi import Ui_Form
import sys
from PyQt5.QtWidgets import QApplication, QWidget, QFileDialog, QMessageBox
from PyQt5.QtCore import Qt, QRect
from Tools.DataStructs import PainterInfo, RectInfo
from loguru import logger
from Tools.QtColorManager import QtColorManager
from Tools.DataStructs import ClassesLabel, DataStructsTools
from Tools.YoloAutoMaker import YoloAutoMaker
import cv2
from YoloImageWidget import YoloImageWidget


class QMainWindow(QWidget, Ui_Form):
    def __init__(self):
        super().__init__()
        self.setupUi(self)
        self.connect_all()
        self.image_list = []
        self.image_index = -1  # 当前索引
        self.label_list: typing.List[ClassesLabel] = []
        self.current_image_path = ''
        self.read_json('/home/hello/graduated_ws/src/yolo_mark_tools/config/label.json')
        self.yolo_auto_marker = YoloAutoMaker(self.label_list,'/home/hello/下载/old-download/yolo11s-seg.pt')

    def connect_all(self):
        self.next_photo_button.clicked.connect(self.on_next_image)
        self.last_photo_button.clicked.connect(self.on_previous_image)
        self.selecet_dir_button.clicked.connect(self.select_dir)
        self.mark_finish_button.clicked.connect(self.on_mark_finish)
        self.select_lable_combox.currentIndexChanged.connect(self.on_select_lable_combox_changed)
        self.select_json_file.clicked.connect(self.read_json)
        self.revoked_button.clicked.connect(self.yolo_image_widget.undo)
        self.redo_button.clicked.connect(self.yolo_image_widget.redo)
        self.auto_mark_button.clicked.connect(self.on_auto_mark_button)
        self.auto_mark_all_button.clicked.connect(self.auto_mark_all)

    def on_auto_mark_button(self):
        rect_list = self.auto_mark(self.current_image_path)
        self.yolo_image_widget.set_rect_list(rect_list)

    def auto_mark(self, image_path):
        cv_image = cv2.imread(image_path)
        print(self.yolo_image_widget.pixmap.width(), self.yolo_image_widget.pixmap.height())
        cv_image = cv2.resize(cv_image, (self.yolo_image_widget.pixmap.width(), self.yolo_image_widget.pixmap.height()))
        rect_list = self.yolo_auto_marker.get_mark_rect(cv_image)
        logger.info(rect_list)
        return rect_list

    def auto_mark_all(self):
        # 从当前文件开始，一直标注到最后
        # todo 目前会阻塞进程，界面会刷新, 但是不能点击
        for image_index in range(self.image_index, len(self.image_list)):
            image_path = self.image_list[image_index]
            rect_list = self.auto_mark(image_path)
            self.mark_finish_one_image(self.get_mark_file_name(image_path),
                                       YoloImageWidget.get_yolo_format_from_rect_list(self.yolo_image_widget,
                                                                                      rect_list))
            self.open_image(image_path)
            self.yolo_image_widget.repaint()
            # self.update()
        QMessageBox.information(self, "提示", "标注完成，共%d张图像." % (len(self.image_list)-self.image_index))
        self.open_image(self.image_list[self.image_index])

    def on_select_lable_combox_changed(self, index: int):
        self.change_painter_info(index)

    def select_dir(self, dir_path: str = ''):
        dir_path = str(dir_path)
        if not os.path.exists(dir_path):
            dir_path = QFileDialog.getExistingDirectory(self, "Select Directory", "",
                                                        QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks)

        if not os.path.exists(dir_path):
            return
        self.image_list = []
        self.image_index = 0
        for f in os.listdir(dir_path):
            if f.endswith('.jpg') or f.endswith('.png'):
                self.image_list.append(os.path.join(dir_path, f))
        self.open_image_by_index(self.image_index)

    def on_next_image(self):
        self.on_mark_finish()
        index = self.image_index + 1
        if index >= len(self.image_list):
            QMessageBox.information(self, "提示", "已经是最后一张.")
            return
        self.image_index += 1
        self.open_image_by_index(index)

    def on_previous_image(self):
        self.on_mark_finish()
        index = self.image_index - 1
        if index < 0:
            QMessageBox.information(self, "提示", "已经是第一张.")
            return
        self.image_index -= 1
        self.open_image_by_index(index)

    def get_mark_file_name(self, image_path):
        return os.path.splitext(image_path)[0] + '.txt'

    def on_mark_finish(self):
        file_name = self.get_mark_file_name(self.current_image_path)
        self.mark_finish_one_image(file_name, self.yolo_image_widget.get_yolo_format())

    def mark_finish_one_image(self, file_name, string):
        logger.info('保存标注信息 to %s' % file_name)
        with open(file_name, 'w') as f:
            f.write(string)

    def open_image_by_index(self, index):
        self.open_image(self.image_list[index])

    def open_image(self, image_path):
        self.current_image_path = image_path

        def file2rect_list(file_path, image_size):
            with open(file_path, 'r') as f:
                lines = f.readlines()
                rect_list = []
                for line in lines:
                    line = line.strip()
                    if len(line) <= 0:
                        continue
                    line_list = line.split(' ')
                    if len(line_list) != 5:
                        continue
                    label_int = int(line_list[0])
                    x_center = float(line_list[1])
                    y_center = float(line_list[2])
                    width = float(line_list[3])
                    height = float(line_list[4])
                    rect = QRect(int(x_center * image_size[0] - width * image_size[0] / 2),
                                 int(y_center * image_size[1] - height * image_size[1] / 2),
                                 int(width * image_size[0]), int(height * image_size[1]))
                    label = self.label_list[label_int]
                    rect_info = RectInfo(rect, label_int,
                                         PainterInfo(QtColorManager().get_color(label_int), label.name))
                    rect_list.append(rect_info)
            return rect_list

        mark_file_name = self.get_mark_file_name(image_path)
        rect_list = []
        try:
            rect_list = file2rect_list(mark_file_name, self.yolo_image_widget.image_size)
        except:
            logger.info('没有相关的标注信息')
        self.yolo_image_widget.openImage(image_path, rect_list)

    def read_json(self, path=''):
        if not os.path.exists(str(path)):
            path, _ = QFileDialog.getOpenFileName(self, "选择json标注文件", "", "Json Files (*.json)")
        if not os.path.exists(path):
            QMessageBox.information(self, "提示", "文件不存在.")
            return
        try:
            self.label_list = DataStructsTools.get_label_list_from_json(path)
        except:
            QMessageBox.information(self, "提示", "读取文件失败.")
            return
        first_label = self.label_list[0]
        init_painter = PainterInfo(QtColorManager().get_color(first_label.index), first_label.name)
        self.yolo_image_widget.set_painter_info(init_painter, first_label.index)
        self.init_combox(self.label_list)

    def change_painter_info(self, index):
        label = self.label_list[index]
        painter_info = PainterInfo(QtColorManager().get_color(label.index), label.name)
        self.yolo_image_widget.set_painter_info(painter_info, label.index)

    def init_combox(self, label_list):
        self.select_lable_combox.clear()
        for label in label_list:
            self.select_lable_combox.addItem('%d-%s' % (label.index, label.name))


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = QMainWindow()
    window.select_dir('/home/hello/桌面/标定/素材视频保存的图片-2')
    window.show()
    sys.exit(app.exec_())
